DiT: Scalable Diffusion Models with Transformers
DiT has a relatively simple idea that it replaces the U-Net in usual Latent Diffusion Models (LDMs) with Transformers (ViT).
Table of Contents
1. Methodology
DiT operates within the LDM framework, using a frozen pretrained VAE to compress images into latent space. The following sections discusses the diffusion model part (i.e., predict noise from latent representation)
1.1. Patchification and Tokenization
Given a latent tensor \(\mathbf{z} \in \mathbb{R}^{I \times I \times C}\), it’s flattened and linearly projected into a sequence of \(T=(\frac{I}{p})^{2}\) tokens of \(d\)-dimensional embedding vectors, where \(p=2,4,8\) denotes the spatial patch size.
Then, a standard fixed 1D or 2D frequency-based (sine-cosine) positional embeddings are added to all patch tokens.
1.2. Conditioning
Diffusion models may require conditioning on timestep \(t\) and conditioning signals (e.g. class label \(c\)). The authors investigate 4 variants, among which the authors recommend adaLN-Zero variant.
1.2.1. In-Context Conditioning
Simply append the vector embeddings of timestep \(t\) and label \(c\) as 2 additional tokens to the input sequence, as if they were image tokens.
1.2.2. Cross-Attention
Concatenate the vector embeddings of \(t,c\) as a sequence of length \(2\), separate from image tokens. Then, modify the Transformer blocks by adding multi-head cross-attention after multi-head self-attention, to make image tokens as \(Q\) and \(t,c\) as \(K,V\).
1.2.3. Adaptive LayerNorm
First, fuse the embeddings of timestep \(t\) and label \(c\), suppose it’s
\[ \mathbf{y} = \mathbf{e}_{t} + \mathbf{e}_{c} \]
In standard LayerNorm, the scale \(\gamma\) and shift \(\beta\) are learnable weights that are fixed after training,
\[ \texttt{LayerNorm}(x) = \gamma \odot \frac{x - \mathbb{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} + \beta \]
But in Adaptive LayerNorm (adaLN), both \(\gamma, \beta\) are predicted by an MLP
\[ [\gamma,\beta] = \texttt{MLP}(y) = \mathrm{Linear}(\mathrm{SiLU}(\mathrm{Linear}(y))) \]
Therefore, adaLN can be expressed as
\[ \mathrm{adaLN}(h,y) = \gamma(y) \odot \left( \frac{h - \mu}{\sigma} \right) + \beta(y) \]
1.2.4. Adaptive LayerNorm Zero
In adaLN-Zero, on top of adaLN, the MLP further predicts a scaler \(\alpha\), such that
\[ h_{l+1} = h_{l} + \alpha_{1} \odot \mathrm{SelfAttn}(\mathrm{adaLN}(h_{l})) \]
So in all, in adaLN-Zero, there are 6 such scalers to be predicted by MLPs: one is adaLN in self-attention, the other is the LayerNorm after self-attention.